import os
import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt

# 定义数据读取函数
def read_image_filenames(data_dir):
    cat_dir = data_dir + 'cat/'
    dog_dir = data_dir + 'dog/'
    # 构建特征数据集，值为对应的图片文件名
    cat_filenames = tf.constant([cat_dir + fn for fn in os.listdir()])
    dog_filenames = tf.constant([dog_dir + fn for fn in os.listdir()])
    filenames = tf.concat([cat_filenames, dog_filenames], axis=-1)
    # 构建标签数据集，cat为0，dog为1
    labels = tf.concat([
        tf.zeros(cat_filenames.shape, dtype=tf.int32),
        tf.ones(dog_filenames.shape, dtype=tf.int32),
    ], axis=-1)
    return filenames, labels

# 定义解码图片和调整图片大小的函数
def decode_image_and_resize(filename, label):
    image_string = tf.io.read_file(filename)
    image_decoded = tf.image.decode_jpeg(image_string)
    # 调整图像大小，要和后面模型输入要求一致，并进行标准化
    image_resized = tf.image.resize(image_decoded, [224, 224]) / 255.0
    return image_resized, label

# ====== 建立猫狗大战Dataset数据集 ======
# train_data_dir = './data/train/'
# filenames, labels = read_image_filenames(train_data_dir)
# dataset = tf.data.Dataset.from_tensor_slices((filenames, labels))
# 取前几个数据
# sub_dataset = dataset.take(3)
# for x, y in sub_dataset:
#     print('filename:', x.numpy(), 'label:', y.numpy())

buffer_size = 20000
batch_size = 8
def prepare_dataset(data_dir):
    filenames, labels = read_image_filenames(data_dir)
    dataset = tf.data.Dataset.from_tensor_slices((filenames, labels))
    dataset = dataset.map(
        map_func=decode_image_and_resize,
        num_parallel_calls=tf.data.experimental.AUTOTUNE
    )
    # 取出前buffer_size个数据放入buffer，并从其中随机采样，采样后的数据用后续数据替换
    dataset = dataset.shuffle(buffer_size)
    dataset = dataset.batch(batch_size)
    dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)
    return dataset

# 读取训练数据并处理
train_data_dir = 'D:\\dogs-vs-cats\\train'
dataset_train = prepare_dataset(train_data_dir)

# 读取数据集，试一下
it = iter(dataset_train)
images, labels = next(it)
print(images.shape)
print(labels.shape)

fig, axs = plt.subplots(1, batch_size)
for i in range(batch_size):
    axs[i].set_title(labels.numpy()[i])
    axs[i].imshow(images.numpy()[i])
    axs[i].set_xticks([])
    axs[i].set_yticks([])
plt.show()
